#ifndef SRC_CRYPTO_HELPER
#define SRC_CRYPTO_HELPER

#include "log_helper.hpp"
#include <NTL/mat_ZZ.h>
#include <helib/helib.h>
#include <omp.h>

#ifdef USE_MPI_WTIME
#include <mpi.h>
#define omp_get_wtime MPI_Wtime
#endif

template <typename Scheme> struct params {};

template <> struct params<helib::CKKS> {
  unsigned long m_lower_bound = 16384;
  unsigned long bits = 119;
  unsigned long precision = 20;
};

template <> struct params<helib::BGV> {
  unsigned long hwt = 120;
  unsigned long m_lower_bound = 20000;
  unsigned long m_upper_bound = 30000;
  unsigned long modulus_chain_bits = 400;
  unsigned long ksm_ncol = 2;
  unsigned long t_bits = 30;
};

struct option_info {
  std::string name;
  void *var;
  bool is_bool;
};

class command_line_parser {
public:
  /**
   * @brief Construct a new `command_line_parser` object
   *        without any option
   */
  command_line_parser();

  /**
   * @brief Construct a new `command_line_parser` object
   *        using BGV/CKKS parameters set
   * @param param
   */
  command_line_parser(params<helib::BGV> &param);
  command_line_parser(params<helib::CKKS> &param);

  // report = true: Feedback command line parsing parameter results
  int parse(int argc, char *argv[], bool report = true);

  /**
   * @brief Add an option `opt` with the string format of its default value
   * `default_str`
   *
   * @param opt             option
   * @param default_str     default value of option `opt`
   * @param name            corresponding variable `name`
   * @param var             Pointer to variable `name`
   */
  void add_option(const std::string &opt, const std::string &default_str,
                  const std::string &name, void *var);
  /**
   * @brief Add an option `opt` without parameters
   *
   * @param opt             option
   * @param name            corresponding variable `name`
   * @param var             Pointer to variable `name`
   */
  void add_option(const std::string &opt, const std::string &name, void *var);

private:
  void update(bool verbose);

private:
  std::unordered_map<std::string, std::string> options_;
  std::unordered_map<std::string, option_info> option_info_;
  void (*print_usage)();
};

class hmm_crypto_agent {
protected:
  hmm_crypto_agent() = default;
  /**
   * @brief Construct a new `hmm_crypto_agent` object
   *        Read from the stream the serialized `Context` object in
   *        binary format
   * @param str   Input `std::istream`
   */
  hmm_crypto_agent(std::istream &str);
  /**
   * @brief Construct a new `hmm_crypto_agent` object
   *        Building the HElib context externally.
   *        * Must provide the building time
   * @param context HElib context
   * @param time    Escaping time of building HElib context
   */
  hmm_crypto_agent(std::shared_ptr<helib::Context> context, double time)
      : context_(context), build_context_time_(time) {}

public:
  virtual ~hmm_crypto_agent() = default;

  std::shared_ptr<helib::Context> context() const { return context_; }
  virtual void report_context(std::ostream &os = std::cout) const;
  virtual void report_time() const = 0;
  double build_context_time() const { return build_context_time_; }

protected:
  std::shared_ptr<helib::Context> context_ = nullptr;

  // * build_context_time (HElib context)
  double build_context_time_ = 0.0;
};

enum ENC { PUBKEY, SECKEY };
template <typename Scheme> class hmm_crypto_client {
private:
  hmm_crypto_client() = delete;

protected:
  /**
   * @brief Construct a new hmm crypto client object
   *        Key Generation using specific `context`
   * @param context HElib context
   */
  hmm_crypto_client(const helib::Context &context);

public:
  virtual ~hmm_crypto_client() { delete sk_; }

  double key_gen_time() const { return key_gen_time_; }

  const helib::PubKey &public_key() const { return *sk_; }

  /**
   * @brief Abbreviation of `helib::Ctxt(hmm_crypto_client::public_key())`
   *
   * @return helib::Ctxt  Empty ciphertext
   */
  helib::Ctxt get_empty_ctxt() const;

  /**
   * @brief Pure virtual encryption method
   *        Must implement its own encoding method
   * @param[out] ctxt   Enc(ptxt)
   * @param[in]  ptxt
   * @param      type   Encrypt using public/secret key (PUBKEY/SECKEY)
   */
  void encrypt(helib::Ctxt &ctxt, const helib::Ptxt<Scheme> &ptxt,
               ENC type = PUBKEY);

  /**
   * @brief Pure virtual encryption method
   *        Must implement its own encoding method
   * @param mat           NTL type plaintext matrix
   * @param type          Encrypt using public/secret key (PUBKEY/SECKEY)
   * @return helib::Ctxt  Enc(mat)
   */
  virtual helib::Ctxt encrypt(const NTL::mat_ZZ &mat, ENC type = PUBKEY) = 0;
  virtual helib::Ctxt encrypt(const std::vector<double> &vec,
                              ENC type = PUBKEY) = 0;
  virtual helib::Ctxt encrypt(const std::vector<long> &vec,
                              ENC type = PUBKEY) = 0;
  /**
   * @brief Pure virtual decryption method
   *        Must implement its own decoding method
   * @param cmat          = Enc(mat)
   * @return NTL::mat_ZZ  = Dec(cmat) = mat
   */
  void decrypt(helib::Ptxt<Scheme> &ptxt, const helib::Ctxt &ctxt);
  virtual void decrypt(NTL::mat_ZZ &mat, const helib::Ctxt &cmat) = 0;
  virtual void decrypt(std::vector<long> &vec, const helib::Ctxt &cmat) = 0;
  virtual void decrypt(std::vector<double> &vec, const helib::Ctxt &cmat) = 0;

protected:
  // NOTE: `SecKey` is derived from `PubKey`
  helib::SecKey *sk_ = nullptr;

  // * key_gen_time (sercet key, public key & key-switching matrices)
  double key_gen_time_ = 0.0;
};

#endif /* SRC_CRYPTO_HELPER */
